import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"  # 原始数据文件路径
OUTPUT_FILE = "monthly_average_data.xlsx"  # 输出文件路径


# ==================== 数据读取与预处理 ====================
def process_data():
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020']  # 仅保留2016-2020年
    all_monthly_data = []

    for year in years:
        df = excel_file.parse(year)
        print(f"\n【{year} 工作表原始数据前5行】\n{df.head()}")

        # 前向填充年、月列（解决空值问题）
        df['年'] = df['年'].ffill()
        df['月'] = df['月'].ffill()

        # 转换为整数（解决浮点问题）
        df['年'] = df['年'].astype(int)
        df['月'] = df['月'].astype(int)

        # 分组计算月均值（跳过含沙量的NaN，pandas.mean()默认忽略NaN）
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'  # mean()会自动跳过NaN
        }).reset_index()

        # 计算排沙量并保留小数
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)

    # 合并并排序数据
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result.to_excel(OUTPUT_FILE, index=False)
    return result


# ==================== 绘制时间序列图 ====================
def plot_time_series(data):
    # 生成时间列（向量化方法，避免apply）
    data['时间'] = pd.to_datetime(data['年'].astype(str) + '-' + data['月'].astype(str) + '-01')

    plt.figure(figsize=(10, 5))
    plt.plot(data['时间'], data['流量(m3/s)'], marker='o', linestyle='-', color='b', label='Monthly Flow')

    # ==================== 修正时间刻度 ====================
    # 生成2014-2022年的年末日期（共9个点：2014-12-31到2022-12-31）
    tick_dates = pd.date_range(start='2014-12-31', end='2022-12-31', freq='1YE')  # 'YE'替代'Y'，避免警告
    # 标签为2014-2022年（共9个，与tick_dates数量一致）
    tick_labels = [str(year) for year in range(2014, 2023)]  # 2014到2022共9年

    plt.xticks(tick_dates, tick_labels, rotation=45)  # 刻度位置与标签数量匹配

    plt.title('Monthly Average Flow (2016-2020)', fontsize=12)
    plt.ylabel('Flow (m³/s)', fontsize=10)
    plt.grid(linestyle='--', alpha=0.7)
    plt.legend()

    plt.savefig('原时间序列图.png', dpi=300, bbox_inches='tight')
    plt.show()


# ==================== 主程序 ====================
if __name__ == "__main__":
    processed_data = process_data()
    filtered_data = processed_data[processed_data['年'].between(2016, 2020)]
    plot_time_series(filtered_data)